package com.yww.datastructure.avltree;

import com.yww.datastructure.stack.LinkedListStack;
import com.yww.datastructure.stack.Stack;

import java.util.function.Consumer;

/**
 * @author yww
 * @description avl树：自平衡二叉树的一种实现方式
 * @since 2024/2/17 16:39
 */
public final class IntAVLTree<V> {
    private AVLNode<V> root;


    public V get(int key) {
        return recurGet(root, key);
    }

    private V recurGet(AVLNode<V> cur, int key) {
        if (cur == null) {
            return null;
        }
        if (key < cur.key) {
            return recurGet(cur.left, key);
        } else if (cur.key < key) {
            return recurGet(cur.right, key);
        } else {
            return cur.val;
        }

    }

    /**
     * 新增方法
     */
    public void put(int key, V val) {
        root = recurPut(root, key, val);
    }

    /**
     * 递归新增方法，新增完成更新高度，调整平衡，返回调整之后节点
     */
    public AVLNode<V> recurPut(AVLNode<V> cur, int key, V val) {
        if (cur == null) {
            return new AVLNode<V>(key, val);
        }
        if (key == cur.key) { // 更新
            cur.val = val;
            return cur;
        }
        if (key < cur.key) {
            cur.left = recurPut(cur.left, key, val);
        } else {
            cur.right = recurPut(cur.right, key, val);
        }
        // 更新高度
        updateHeight(cur);
        // 调整平衡
        return balance(cur);
    }

    /**
     * 递归进行删除
     */
    public void remove(int key) {
        root = recurRemove(root, key);
    }

    private AVLNode<V> recurRemove(AVLNode<V> cur, int key) {
        if (cur == null) {
            return null;
        }
        if (key < cur.key) {
            cur.left = recurRemove(cur.left, key);
        } else if (cur.key < key) {
            cur.right = recurRemove(cur.right, key);
        } else { // 找到了
            if (cur.left == null && cur.right == null) { // 情况1：没有孩子
                return null;
            } else if (cur.left == null) { // 情况2：只有右孩子
                cur = cur.right;
            } else if (cur.right == null) { // 情况3：只有左孩子
                cur = cur.left;
            } else { // 情况4：左右孩子都有
                AVLNode<V> s = cur.right; // 后继
                while (s.left != null) {
                    s = s.left;
                }
                s.right = recurRemove(cur.right, s.key); // 先更新右
                s.left = cur.left; // 再更新左
                cur = s;
            }
        }
        updateHeight(cur); // 调整高度
        return balance(cur); // 调整平衡

    }

    public void inorder(Consumer<V> consumer) {
        Stack<AVLNode<V>> stack = new LinkedListStack<>();
        AVLNode<V> cur = root;
        while (cur != null || !stack.isEmpty()) {
            if (cur != null) {
                stack.push(cur);
                cur = cur.left;
            } else {
                AVLNode<V> pop = stack.pop();
                if (consumer != null) {
                    consumer.accept(pop.val);
                }
                cur = pop.right;
            }
        }
    }


    /**
     * 获取树高
     */
    private int height(AVLNode<V> node) {
        if (node == null) {
            return 0;
        }
        return node.height;
    }

    /**
     * 更新树高
     */
    private void updateHeight(AVLNode<V> node) {
        node.height = Integer.max(height(node.left), height(node.right)) + 1;
    }

    /**
     * 获取一颗树的平衡因子（balance Factor）<br>
     * 平衡因子 = 左子树高度 - 右子树高度 <br>
     * bf > 1 左子树太高 <br>
     * bf < -1 右子树太高 <br>
     *
     * @return -1 0 1 树达到平衡 <br>
     */
    private int bf(AVLNode<V> node) {
        return height(node.left) - height(node.right);
    }

    /**
     * 对node失衡节点进行右旋（左子树太高 LL）
     *
     * @param red 需要进行右旋的失衡节点
     * @return 旋转完成的根节点
     */
    private AVLNode<V> rightRotate(AVLNode<V> red) {
        AVLNode<V> yellow = red.left;
        AVLNode<V> green = yellow.right;
        // 红下沉，黄上位
        yellow.right = red;
        red.left = green;
        // 更新高度
        updateHeight(red);
        updateHeight(yellow);
        return yellow;
    }

    /**
     * 对node失衡节点进行左旋（右子树太高 RR）
     *
     * @param red 需要进行左旋的失衡节点
     * @return 旋转完成的根节点
     */

    private AVLNode<V> leftRotate(AVLNode<V> red) {
        AVLNode<V> yellow = red.right;
        AVLNode<V> green = yellow.left;
        // 红下沉，黄上位
        yellow.left = red;
        red.right = green;
        // 更新高度（先更新下沉的红高度，在更新上位黄的高度）
        updateHeight(red);
        updateHeight(yellow);
        return yellow;
    }

    /**
     * 对red失衡节点进行左右旋（red的左子树yellow太高，yellow的右子树太高 LR）<br>
     * 先对yellow的右子树进行左旋，然后对red进行右旋 <br>
     *
     * @param red 需要进行左右的失衡节点
     * @return 旋转完成的根节点
     */
    private AVLNode<V> leftRightRotate(AVLNode<V> red) {
        red.left = leftRotate(red.left);
        return rightRotate(red);
    }

    /**
     * 对red失衡节点进行右左旋（red的右子树yellow太高，yellow的左子树太高 RL）<br>
     * 先对yellow的左子树进行右旋，然后对red进行左旋 <br>
     *
     * @param red 需要进行右左的失衡节点
     * @return 旋转完成的根节点
     */
    private AVLNode<V> rightLeftRotate(AVLNode<V> red) {
        red.right = rightRotate(red.right);
        return leftRotate(red);
    }

    /**
     * 调整node的平衡，返回调整之后的根节点
     */
    private AVLNode<V> balance(AVLNode<V> node) {
        if (node == null) {
            return null;
        }
        int bf = bf(node);
        if (bf > 1 && bf(node.left) >= 0) { // LL
            return rightRotate(node); // 右旋
        } else if (bf > 1 && bf(node.left) < 0) { // LR
            return leftRightRotate(node); // 左右旋
        } else if (bf < -1 && bf(node.right) > 0) { // RL
            return rightLeftRotate(node); // 右左旋
        } else if (bf < -1 && bf(node.right) <= 0) { // RR
            return leftRotate(node); // 左旋
        }
        return node;
    }

    /**
     * 节点定义
     */
    private static class AVLNode<V> {
        private int height = 1; // 树高度
        private final int key;
        private V val;
        private AVLNode<V> left;
        private AVLNode<V> right;

        public AVLNode(int key, V val) {
            this.key = key;
            this.val = val;
        }
    }
}
